Skip to content

[Offloading] [Bugfix] Fix disk offloading of models with explicit tensor dtypes#46849

Merged
SunMarc merged 2 commits into
huggingface:mainfrom
kylesayrs:kylesayrs/bugfix-disk-dtype
Jun 24, 2026
Merged

[Offloading] [Bugfix] Fix disk offloading of models with explicit tensor dtypes#46849
SunMarc merged 2 commits into
huggingface:mainfrom
kylesayrs:kylesayrs/bugfix-disk-dtype

Conversation

@kylesayrs

@kylesayrs kylesayrs commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Purpose

  • Fix bug where DSV4 would load tid2eid with the incorrect precision, leading to incorrect hash gate values
    • Casting from int64 to bfloat16 was causing tid2eid to round to out-of-bounds values, not to mention that the values were incorrect

Changes

  • Instead of using model.dtype to determine the load dtype for disk-offloaded weights, instead use the dtype of the associated meta tensor in the meta model. The meta tensor dtype is already the correct dtype, influenced by the local_torch_dtype load context
    • This dtype field is read by OffloadedWeightsLoader here

Testing

  • Added test_disk_onload_dtype
RUN_SLOW=1 python3 -m pytest tests/utils/test_modeling_utils.py::ModelUtilsTest::test_disk_onload_dtype -sx

Suggested Reviewers

@SunMarc @zucchini-nlp

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
@kylesayrs kylesayrs changed the title [Offloading] [Bugfix] Use meta tensor dtype when onloading from disk [Offloading] [Bugfix] Fix disk offloading of models with explicit tensor dtypes Jun 23, 2026
@github-actions

Copy link
Copy Markdown
Contributor

CI Dashboard: View test results in Grafana

@SunMarc SunMarc left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, looks good to me !

@SunMarc SunMarc enabled auto-merge June 24, 2026 14:02
@SunMarc SunMarc added this pull request to the merge queue Jun 24, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jun 24, 2026
@SunMarc SunMarc added this pull request to the merge queue Jun 24, 2026
Merged via the queue into huggingface:main with commit 176a17b Jun 24, 2026
104 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants